package com.api.chatai.utils;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.knuddels.jtokkit.Encodings;
import com.knuddels.jtokkit.api.Encoding;
import com.knuddels.jtokkit.api.EncodingRegistry;
import com.knuddels.jtokkit.api.EncodingType;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.json.JSONArray;
import org.json.JSONObject;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.concurrent.TimeUnit;

/**
 * 流式传输请求封装
 */
@Slf4j
public class StreamChatUtil {


    public String chat(String message, HttpServletResponse res) {

        res.setContentType("text/event-stream");
        res.setCharacterEncoding("UTF-8");


//      String api = "https://api.xty.app/v1/chat/completions";
        String api = "https://hk.xty.app/v1/chat/completions";
        String apiKey = "sk-e8bOt4XipXTm68f117790a690cA141729d514eF3EbEfF966";


        // 封装json
        JSONObject requestBody = new JSONObject();
        requestBody.put("model", "gpt-3.5-turbo");
        requestBody.put("stream", true);
        JSONArray messagesArray = new JSONArray();

        //系统
        JSONObject messageObjectSystem = new JSONObject();
        messageObjectSystem.put("role", "system");
        messageObjectSystem.put("content", "你是一个全能助手，能回答各种问题。");
        messagesArray.put(messageObjectSystem);

        //用户
        JSONObject messageObjectUser = new JSONObject();
        messageObjectUser.put("role", "user");
        messageObjectUser.put("content", message);
        messagesArray.put(messageObjectUser);

        requestBody.put("messages", messagesArray);

        log.info("请求体: {} \n", requestBody);


        Request request = new Request.Builder()
                .url(api)
                .post(RequestBody.create(requestBody.toString(), MediaType.get("application/json")))
                .addHeader("Authorization", "Bearer " + apiKey)
                .build();

        // 发请求
        OkHttpClient client = new OkHttpClient.Builder().followRedirects(true)
                .connectTimeout(30, TimeUnit.SECONDS)
                .callTimeout(300, TimeUnit.SECONDS)
                .pingInterval(5, TimeUnit.SECONDS)
                .readTimeout(120, TimeUnit.SECONDS)
                .writeTimeout(120, TimeUnit.SECONDS)
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
//               System.out.println("请求失败: " + response);
                log.error("请求失败，请求体: {} \n", requestBody);
            }


            StringBuilder allContent = new StringBuilder();
            ObjectMapper mapper = new ObjectMapper();
            try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.body().byteStream()))) {
                String line;
                while ((line = reader.readLine()) != null) {

//               log.info("响应体: {} \n", line);
                    int startIndex = line.indexOf('{');
                    if (startIndex != -1) {
                        ObjectNode dataJson = (ObjectNode) mapper.readTree(line.substring(startIndex));
                        if (!dataJson.has("choices") || dataJson.withArray("choices").isEmpty()) continue;

                        ObjectNode choicesJson = (ObjectNode) dataJson.withArray("choices").get(0);
                        ObjectNode deltaJson = (ObjectNode) choicesJson.get("delta");

                        if (deltaJson != null && deltaJson.has("content")) {
                            String content = deltaJson.get("content").asText();
                            allContent.append(content); // 将接收到的内容累加到 allContent 变量中,公式中的\要转义
                            res.getWriter().write("data: " + content.replace("\n", "\\x0A") + "\n\n");
                            res.getWriter().flush();
                        }
                    }
                }
                res.getWriter().write("data: [DONE]\n\n");
                res.getWriter().flush();

//            System.out.println(allContent);
//            System.out.println("共使用 : " + totalTokensUsed + " token");
//            log.info("回复内容 : {}", allContent);
                return allContent.toString();

            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        return null;
    }

    //模拟回答
    public String mockChat(String message, HttpServletResponse res) {
        res.setContentType("text/event-stream");
        res.setCharacterEncoding("UTF-8");


        StringBuilder allContent = new StringBuilder();
        try {
            //循环100次,随机打印英文单词和中文单词
            for (int i = 0; i < 10; i++) {
                String content;
                if (i % 2 == 0) {
                    content = "你好";
                    allContent.append(content);
                } else {
                    content = "hello";
                    allContent.append(content);
                }
                res.getWriter().write("data: " + content);
                res.getWriter().flush();
                try {
                    Thread.sleep(500);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }

            res.getWriter().write("data: [DONE]\n\n");
            res.getWriter().flush();
            return allContent.toString();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;

    }


}

